package com.shujia.util

import org.apache.spark.internal.Logging
import org.apache.spark.sql.expressions.UserDefinedFunction
import org.apache.spark.sql.functions.udf
import org.apache.spark.sql.{DataFrame, SaveMode, SparkSession}

abstract class SparkTool extends Logging {

  /**
    * 子类会继承父类的mainH函数
    *
    */

  var spark: SparkSession = _
  //天分区
  var day_id: String = _
  //月分区
  var month_id: String = _


  //编写自定义函数，实现加密处理
  val hashMd5: UserDefinedFunction = udf((str: String) => {
    Md5.hashMD5(str + "shujia")
  })


  def main(args: Array[String]): Unit = {

    /**
      * 获取时间参数
      *
      */

    if (args.length < 1) {
      log.error("请指定时间参数：day_id")
      return
    }

    val par = args(0)

    //月分区
    if (par.length == 6) {
      month_id = par
      log.info(s"正在处理的分区是：$month_id")
    }


    //天分区
    if (par.length == 8) {
      //天分区
      day_id = args(0)
      //月分区
      month_id = day_id.substring(0, 6)
      log.info(s"正在处理的分区是：$day_id")
      log.info(s"月分区：$month_id")
    }




    //创建spark环境
    spark = SparkSession.builder()
      .appName(this.getClass.getSimpleName.replace("$", ""))
      .enableHiveSupport() //开启hive的元数据支持
      .getOrCreate()




    //如果在sql厚葬要使用自定义函数需要注册
    spark.udf.register("hashMd5", hashMd5)


    /**
      * 调用子类的方法
      *
      */
    this.run(spark)


  }

  /**
    * 在子类中实现run方法，实现自定义的代码逻辑
    *
    * 在子类run方法的前面加上下面两行代码
    * import spark.implicits._
    * import org.apache.spark.sql.functions._
    *
    * @param spark spark的环境
    */
  def run(spark: SparkSession)


  /**
    * 保存数据和增加天分区
    *
    * @param df      数据
    * @param path    保存的路径
    * @param dbTable 库和表
    * @param sep     数据分隔方式
    * @param format  数据格式
    */
  def saveDataAndAddPartitionWithDay(df: DataFrame, path: String, dbTable: String, sep: String = "\t", format: String = "csv"): Unit = {

    df
      .write
      .format(format)
      .option("sep", sep)
      .mode(SaveMode.Overwrite)
      .save(s"$path/day_id=$day_id")

    /**
      * 增加分区
      *
      */

    spark.sql(
      s"""
         |alter table $dbTable  add if not exists partition(day_id=$day_id)
         |
      """.stripMargin)


  }


  /**
    * 保存数据和增加月分区
    *
    * @param df      数据
    * @param path    保存的路径
    * @param dbTable 库和表
    * @param sep     数据分隔方式
    * @param format  数据格式
    */
  def saveDataAndAddPartitionWithMonth(df: DataFrame, path: String, dbTable: String, sep: String = "\t", format: String = "csv"): Unit = {

    df
      .write
      .format(format)
      .option("sep", sep)
      .mode(SaveMode.Overwrite)
      .save(s"$path/month_id=$month_id")

    /**
      * 增加分区
      *
      */

    spark.sql(
      s"""
         |alter table $dbTable  add if not exists partition(month_id=$month_id)
         |
      """.stripMargin)


  }

}
